from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_ollama import ChatOllama
from dotenv import load_dotenv
import os


class MultiDialogerWithMemory:
    """带记忆的多轮对话类"""
    
    def __init__(self):
        # 语言模型实例
        self.llm = None
        # 带消息历史记录的可运行对象
        self.with_message_history = None
        # 存储不同会话的历史记录的字典
        self.store = {}
        
    def init_model(self):
        """初始化模型"""
        self.llm = ChatOllama(model="deepseek-r1:14b", validate_model_on_init=True)
        print("模型初始化完成！")
        
    def create_memory_model(self):
        """创建带记忆的模型"""
        if self.llm is None:
            raise ValueError("请先调用 init_model() 初始化模型")
            
        # 创建提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个友好、有帮助的AI助手。请用中文回答问题。"),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{input}")
        ])
        
        # 创建链
        chain = prompt | self.llm
        
        # 创建带历史记录的可运行对象
        self.with_message_history = RunnableWithMessageHistory(
            chain,
            self._get_session_history,
            input_messages_key="input",
            history_messages_key="history",
        )
        print("带记忆的模型创建完成！")
        
    def _get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
        """获取会话历史记录"""
        if session_id not in self.store:
            self.store[session_id] = InMemoryChatMessageHistory()
        return self.store[session_id]
        
    def chat(self, session_id: str = "default_session"):
        """开始多轮对话"""
        if self.with_message_history is None:
            raise ValueError("请先调用 create_memory_model() 创建带记忆的模型")
            
        print("=== 多轮对话开始 ===")
        print("提示：输入 'quit' 结束对话")
        print("-" * 50)
        
        while True:
            try:
                # 获取用户输入
                user_input = input("\n你: ").strip()
                
                # 检查是否退出
                if user_input.lower() == 'quit':
                    print("\n对话结束，再见！")
                    break
                    
                # 如果输入为空，继续循环
                if not user_input:
                    print("请输入有效内容...")
                    continue
                    
                # 调用模型获取回复
                response = self.with_message_history.invoke(
                    {"input": user_input},
                    config={"configurable": {"session_id": session_id}}
                )
                
                # 显示AI回复
                print(f"\nAI: {response.content}")
                
            except KeyboardInterrupt:
                print("\n\n对话被中断，再见！")
                break
            except Exception as e:
                print(f"\n发生错误: {e}")
                print("请重试...")
                
    def get_chat_history(self, session_id: str = "default_session"):
        """获取聊天历史记录"""
        if session_id in self.store:
            history = self.store[session_id]
            print(f"\n=== 会话 {session_id} 的历史记录 ===")
            for message in history.messages:
                role = "用户" if message.type == "human" else "AI"
                print(f"{role}: {message.content}")
        else:
            print(f"会话 {session_id} 没有历史记录")
            
    def clear_history(self, session_id: str = "default_session"):
        """清除指定会话的历史记录"""
        if session_id in self.store:
            del self.store[session_id]
            print(f"会话 {session_id} 的历史记录已清除")
        else:
            print(f"会话 {session_id} 没有历史记录")


def main():
    """主函数演示"""
    # 创建对话器实例
    dialoger = MultiDialogerWithMemory()
    
    # 初始化模型
    dialoger.init_model()
    
    # 创建带记忆的模型
    dialoger.create_memory_model()
    
    # 开始聊天
    dialoger.chat()
    
    # 可选：查看聊天历史
    # dialoger.get_chat_history()


if __name__ == "__main__":
    main()